{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e916a938",
   "metadata": {},
   "source": [
    "## Index\n",
    "\n",
    "* [Introduction](#intro)\n",
    "* [Preparation](#preparation)\n",
    "* [Optimization](#optimization)\n",
    "    * [Normalization](#normalization)\n",
    "    * [Object detection](#object-detection)\n",
    "    * [Dimensionality Reduction](#dimensionality-reduction)\n",
    "* [Online Demo](#demo)\n",
    "\n",
    "# Deep Dive into Real-World Image Search Engine with Towhee <a class=\"anchor\" id=\"intro\"></a>\n",
    "\n",
    "In the [previous tutorial](./1_build_image_search_engine.ipynb), we built and prototyped a proof-of-concept image search engine. With test results from the previous tutorial, we find out that more complex model usually generates larger embeddings, hence leads to better search performance but slower speed.\n",
    "\n",
    "Now, let's try some new methods to improve performance and save resource, other than changing model. At the end, we will also learn how to deploy it as a simple online demo. With this tutorial, you are able to build a reverse image search engine more practical in production."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4f5487d",
   "metadata": {},
   "source": [
    "## Preparation <a class=\"anchor\" id=\"preparation\"></a>\n",
    "\n",
    "Here is a table of search performance with different models from the previous tutorial. We will make some improvement in pipelines and compare model performance in this tutorial. Before getting started, we need to prepare dependencies, example data, and helpful functions, which have detailed explanation in the previous tutorial.\n",
    "\n",
    "| model | dim | mAP@10 | qps |\n",
    "| -- | -- | -- | -- |\n",
    "| vgg16 | 512 | 0.658 | 53 |\n",
    "| resnet50 | 2048 | 0.886 | 35 |\n",
    "| tf_efficientnet_b7 | 2560 | 0.983 | 16 |\n",
    "\n",
    "**Install dependencies**: install python dependencies with proper versions for your environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "35e97f15",
   "metadata": {},
   "outputs": [],
   "source": [
    "! python -m pip -q install towhee gradio==3.3 opencv-python"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0d15992",
   "metadata": {},
   "source": [
    "**Prepare data**: download example data, which is a subset of [ImageNet](https://www.image-net.org/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e9247127",
   "metadata": {},
   "outputs": [],
   "source": [
    "! curl -L https://github.com/towhee-io/examples/releases/download/data/reverse_image_search.zip -O\n",
    "! unzip -q -o reverse_image_search.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a497d876",
   "metadata": {},
   "source": [
    "**Start Milvus:** install and start Milvus service.\n",
    "\n",
    "This notebook uses [milvus 2.2.10](https://milvus.io/docs/v2.2.x/install_standalone-docker.md) and [pymilvus 2.2.11](https://milvus.io/docs/release_notes.md#2210)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "83eaaac8",
   "metadata": {},
   "outputs": [],
   "source": [
    "! wget https://github.com/milvus-io/milvus/releases/download/v2.2.10/milvus-standalone-docker-compose.yml -O docker-compose.yml\n",
    "! docker-compose up -d\n",
    "! python -m pip install -q pymilvus==2.2.11"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acd8edad",
   "metadata": {},
   "source": [
    "**Helpful functions**: import necessary packages, set parameters, and build helpful functions in advance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "58725efe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy\n",
    "import time\n",
    "import csv\n",
    "from glob import glob\n",
    "from pathlib import Path\n",
    "from statistics import mean\n",
    "\n",
    "from towhee import pipe, ops, DataCollection\n",
    "from towhee.types.image import Image\n",
    "from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility\n",
    "\n",
    "# Towhee parameters\n",
    "MODEL = 'vgg16'\n",
    "DEVICE = None # if None, use default device (cuda is enabled if available)\n",
    "\n",
    "# Milvus parameters\n",
    "HOST = '127.0.0.1'\n",
    "PORT = '19530'\n",
    "TOPK = 10\n",
    "DIM = 512 # dimension of embedding extracted, change with MODEL\n",
    "COLLECTION_NAME = 'deep_dive_image_search_' + MODEL\n",
    "INDEX_TYPE = 'IVF_FLAT'\n",
    "METRIC_TYPE = 'L2'\n",
    "\n",
    "# patterns of image paths\n",
    "INSERT_SRC = './train/*/*.JPEG'\n",
    "QUERY_SRC = './test/*/*.JPEG'\n",
    "\n",
    "to_insert = glob(INSERT_SRC)\n",
    "to_test = glob(QUERY_SRC)\n",
    "\n",
    "# Create milvus collection (delete first if exists)\n",
    "def create_milvus_collection(collection_name, dim):\n",
    "    if utility.has_collection(collection_name):\n",
    "        utility.drop_collection(collection_name)\n",
    "    \n",
    "    fields = [\n",
    "        FieldSchema(name='path', dtype=DataType.VARCHAR, description='path to image', max_length=500, \n",
    "                    is_primary=True, auto_id=False),\n",
    "        FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, description='image embedding vectors', dim=dim)\n",
    "    ]\n",
    "    schema = CollectionSchema(fields=fields, description='reverse image search')\n",
    "    collection = Collection(name=collection_name, schema=schema)\n",
    "\n",
    "    index_params = {\n",
    "        'metric_type': METRIC_TYPE,\n",
    "        'index_type': INDEX_TYPE,\n",
    "        'params': {\"nlist\": 2048}\n",
    "    }\n",
    "    collection.create_index(field_name='embedding', index_params=index_params)\n",
    "    return collection\n",
    "            \n",
    "# Read images\n",
    "decoder = ops.image_decode('rgb').get_op()\n",
    "def read_images(img_paths):\n",
    "    imgs = []\n",
    "    for p in img_paths:\n",
    "        img = decoder(p)\n",
    "        imgs.append(img)\n",
    "#         imgs.append(Image(cv2.imread(p), 'RGB'))\n",
    "    return imgs\n",
    "\n",
    "# Get ground truth\n",
    "def ground_truth(path):\n",
    "    train_path = str(Path(path).parent).replace('test', 'train')\n",
    "    return [str(Path(x).resolve()) for x in glob(train_path + '/*.JPEG')]\n",
    "\n",
    "# Calculate Average Precision\n",
    "def get_ap(pred: list, gt: list):\n",
    "    ct = 0\n",
    "    score = 0.\n",
    "    for i, n in enumerate(pred):\n",
    "        if n in gt:\n",
    "            ct += 1\n",
    "            score += (ct / (i + 1))\n",
    "    if ct == 0:\n",
    "        ap = 0\n",
    "    else:\n",
    "        ap = score / ct\n",
    "    return ap"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1aadd8e8",
   "metadata": {},
   "source": [
    "## Optimization <a class=\"anchor\" id=\"optimization\"></a>\n",
    "\n",
    "In the previous tutorial, we have measured the search performance with **mAP** and compared performance for different models. This tutorial will show how to improve performance by normalization, implement pipeline with object detection, and reduce dimension to save resource. \n",
    "\n",
    "### Normalization <a class=\"anchor\" id=\"normalization\"></a>\n",
    "\n",
    "A quick optimization is normalizing the embedding features before indexing them in Milvus. Thus, the L2 metric used by Milvus is equivalent to cosine similarity, which measures the similarity using the angle between vectors while ignoring the magnitude of vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "54480275",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embedding pipeline\n",
    "p_embed = (\n",
    "    pipe.input('img_path')\n",
    "        .map('img_path', 'img', ops.image_decode('rgb'))\n",
    "        .map('img', 'vec', ops.image_embedding.timm(model_name=MODEL, device=DEVICE))\n",
    "        .map('vec', 'vec', lambda x: x / numpy.linalg.norm(x, axis=0))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1c7f5ecd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table style=\"border-collapse: collapse;\"><tr><th style=\"text-align: center; font-size: 130%; border: none;\">img_path</th> <th style=\"text-align: center; font-size: 130%; border: none;\">img</th> <th style=\"text-align: center; font-size: 130%; border: none;\">vec</th></tr> <tr><td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \">./train/parachute/n03888257_31784.JPEG</td> <td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td> <td style=\"text-align: left; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \">[0.018488767, 0.069933996, 0.00011331425, ...] shape=(512,)</td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Display embedding result, no need for implementation\n",
    "p_display = p_embed.output('img_path', 'img', 'vec')\n",
    "\n",
    "DataCollection(p_display(to_insert[0])).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00c3993f",
   "metadata": {},
   "source": [
    "Now we have an embedding pipeline extracting normalized vectors for images. Let's build a image search engine based on the embedding pipeline and Milvus collection. We evaluate the engine by inserting candidate data and querying test images. The result table below shows mAP increases for all models. This proves that normalization is able to improve image search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9bfeb0a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A new collection created: deep_dive_image_search_vgg16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "E0209 17:36:01.144905126   63043 fork_posix.cc:76]           Other threads are currently calling into gRPC, skipping fork() handlers\n",
      "E0209 17:36:01.296288564   60015 fork_posix.cc:76]           Other threads are currently calling into gRPC, skipping fork() handlers\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of data inserted: 1000\n",
      "mAP@10: 0.7381725828294281\n"
     ]
    }
   ],
   "source": [
    "# Connect to Milvus service\n",
    "connections.connect(host=HOST, port=PORT)\n",
    "\n",
    "# Create collection\n",
    "collection = create_milvus_collection(COLLECTION_NAME, DIM)\n",
    "print(f'A new collection created: {COLLECTION_NAME}')\n",
    "\n",
    "# Insert data\n",
    "p_insert = (\n",
    "        p_embed.map(('img_path', 'vec'), 'mr', ops.ann_insert.milvus_client(\n",
    "                    host=HOST,\n",
    "                    port=PORT,\n",
    "                    collection_name=COLLECTION_NAME\n",
    "                    ))\n",
    "          .output('mr')\n",
    ")\n",
    "\n",
    "for img_path in to_insert:\n",
    "    p_insert(img_path)\n",
    "print('Number of data inserted:', collection.num_entities)\n",
    "\n",
    "# Performance\n",
    "collection.load()\n",
    "p_search_pre = (\n",
    "        p_embed.map('vec', ('search_res'), ops.ann_search.milvus_client(\n",
    "                    host=HOST, port=PORT, limit=TOPK,\n",
    "                    collection_name=COLLECTION_NAME))\n",
    "               .map('search_res', 'pred', lambda x: [str(Path(y[0]).resolve()) for y in x])\n",
    "#                .output('img_path', 'pred')\n",
    ")\n",
    "p_eval = (\n",
    "    p_search_pre.map('img_path', 'gt', ground_truth)\n",
    "                .map(('pred', 'gt'), 'ap', get_ap)\n",
    "                .output('ap')\n",
    ")\n",
    "\n",
    "res = []\n",
    "for img_path in to_test:\n",
    "    ap = p_eval(img_path).get()[0]\n",
    "    res.append(ap)\n",
    "\n",
    "mAP = mean(res)\n",
    "\n",
    "print(f'mAP@{TOPK}: {mAP}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79294a06",
   "metadata": {},
   "source": [
    "| model | mAP@10 (no norm) | mAP@10 (norm) |\n",
    "| -- | -- | -- |\n",
    "| vgg16 | 0.658 | 0.738 |\n",
    "| resnet50 | 0.886 | 0.917 |\n",
    "| tf_efficientnet_b7 | 0.983 | 0.988 |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f1d1562",
   "metadata": {},
   "source": [
    "###  Object Detection <a class=\"anchor\" id=\"object-detection\"></a>\n",
    "\n",
    "Another common option in reverse image search is object detection. Sometimes the search engine is distracted by small objects or background in the image. Cropping the original image and querying only the main object help to resolve this issue.\n",
    "\n",
    "Let's take a look at a bad search. With normalized embeddings extracted by `vgg16`, the test image *'./test/rocking_chair/n04099969_23803.JPEG'* gets a list of similar images containing some incorrect results, which has an Average Precision of 0.347."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e9ddee4c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "E0209 17:36:52.111067087   67628 fork_posix.cc:76]           Other threads are currently calling into gRPC, skipping fork() handlers\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table style=\"border-collapse: collapse;\"><tr><th style=\"text-align: center; font-size: 130%; border: none;\">img_path</th> <th style=\"text-align: center; font-size: 130%; border: none;\">img</th> <th style=\"text-align: center; font-size: 130%; border: none;\">res</th> <th style=\"text-align: center; font-size: 130%; border: none;\">ap</th></tr> <tr><td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \">./test/rocking_chair/n04099969_23803.JPEG</td> <td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td> <td style=\"text-align: center; vertical-align: top; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td> <td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \">0.34722222222222215</td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "p_search_img = (\n",
    "    p_search_pre.map('img_path', 'gt', ground_truth)\n",
    "                .map(('pred', 'gt'), 'ap', get_ap)\n",
    "                .map('pred', 'res', read_images)\n",
    "                .output('img_path', 'img', 'res', 'ap')\n",
    ")\n",
    "DataCollection(p_search_img('./test/rocking_chair/n04099969_23803.JPEG')).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ce01146",
   "metadata": {},
   "source": [
    "Now let's preprocess the test image by focusing on the main object in it. Here we use YOLOv5 to get objects in the image. We select the object with the largest area in the original image, and then search across database with the object image.\n",
    "\n",
    "- `get_object`: a function to get the image of the largest object detecte, or the original imageif there is no object\n",
    "- `p_yolo`: a pipeline to crop the largest object in the given image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b2590a56",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_max_object(img, boxes):\n",
    "    if len(boxes) == 0:\n",
    "        return img\n",
    "    max_area = 0\n",
    "    for box in boxes:\n",
    "        x1, y1, x2, y2 = box\n",
    "        area = (x2-x1)*(y2-y1)\n",
    "        if area > max_area:\n",
    "            max_area = area\n",
    "            max_img = img[y1:y2,x1:x2,:]\n",
    "    return max_img\n",
    "\n",
    "p_yolo = (\n",
    "    pipe.input('img_path')\n",
    "        .map('img_path', 'img', ops.image_decode('rgb'))\n",
    "        .map('img', ('boxes', 'class', 'score'), ops.object_detection.yolov5())\n",
    "        .map(('img', 'boxes'), 'object', get_max_object)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "845d8504",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pip.repos.neuron.amazonaws.com\n",
      "Requirement already satisfied: pyyaml in /home/mengjia.gu/anaconda3/lib/python3.9/site-packages (6.0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\n",
      "Using cache found in /home/data3/cache/torch/hub/ultralytics_yolov5_master\n",
      "YOLOv5 🚀 2023-1-30 Python-3.9.12 torch-1.13.1+cu117 CUDA:0 (NVIDIA GeForce RTX 3080 Ti, 12054MiB)\n",
      "\n",
      "Fusing layers... \n",
      "YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients\n",
      "Adding AutoShape... \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table style=\"border-collapse: collapse;\"><tr><th style=\"text-align: center; font-size: 130%; border: none;\">img</th> <th style=\"text-align: center; font-size: 130%; border: none;\">object</th></tr> <tr><td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td> <td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Display embedding result, no need for implementation\n",
    "p_display = (\n",
    "    p_yolo.output('img', 'object')\n",
    ")\n",
    "DataCollection(p_display('./test/rocking_chair/n04099969_23803.JPEG')).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4afcf448",
   "metadata": {},
   "source": [
    "With object detection, we search for *'./test/rocking_chair/n04099969_23803.JPEG'* again across the same Milvus collection. The average precision has increased by about 45%. It is a great improvement for this query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2bff57cb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pip.repos.neuron.amazonaws.com\n",
      "Requirement already satisfied: pyyaml in /home/mengjia.gu/anaconda3/lib/python3.9/site-packages (6.0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\n",
      "Using cache found in /home/data3/cache/torch/hub/ultralytics_yolov5_master\n",
      "YOLOv5 🚀 2023-1-30 Python-3.9.12 torch-1.13.1+cu117 CUDA:0 (NVIDIA GeForce RTX 3080 Ti, 12054MiB)\n",
      "\n",
      "Fusing layers... \n",
      "YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients\n",
      "Adding AutoShape... \n",
      "E0209 17:37:09.336641703   60015 fork_posix.cc:76]           Other threads are currently calling into gRPC, skipping fork() handlers\n",
      "E0209 17:37:09.455251735   68535 fork_posix.cc:76]           Other threads are currently calling into gRPC, skipping fork() handlers\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table style=\"border-collapse: collapse;\"><tr><th style=\"text-align: center; font-size: 130%; border: none;\">img</th> <th style=\"text-align: center; font-size: 130%; border: none;\">object</th> <th style=\"text-align: center; font-size: 130%; border: none;\">res</th> <th style=\"text-align: center; font-size: 130%; border: none;\">ap</th></tr> <tr><td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td> <td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td> <td style=\"text-align: center; vertical-align: top; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td> <td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \">0.775</td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Search\n",
    "p_search_pre_yolo = (\n",
    "    p_yolo.map('object', 'vec', ops.image_embedding.timm(model_name=MODEL, device=DEVICE))\n",
    "          .map('vec', 'vec', lambda x: x / numpy.linalg.norm(x, axis=0))\n",
    "          .map('vec', ('search_res'), ops.ann_search.milvus_client(\n",
    "                host=HOST, port=PORT, limit=TOPK,\n",
    "                collection_name=COLLECTION_NAME))\n",
    "          .map('search_res', 'pred', lambda x: [str(Path(y[0]).resolve()) for y in x])\n",
    "#          .output('img_path', 'pred')\n",
    ")\n",
    "\n",
    "# Evaluate with AP\n",
    "p_search_img_yolo = (\n",
    "    p_search_pre_yolo.map('img_path', 'gt', ground_truth)\n",
    "                     .map(('pred', 'gt'), 'ap', get_ap)\n",
    "                     .map('pred', 'res', read_images)\n",
    "                     .output('img', 'object', 'res', 'ap')\n",
    ")\n",
    "DataCollection(p_search_img_yolo('./test/rocking_chair/n04099969_23803.JPEG')).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40eec2ce",
   "metadata": {},
   "source": [
    "### Dimensionality Reduction <a class=\"anchor\" id=\"dimensionality-reduction\"></a>\n",
    "\n",
    "For a system in production, it is practical to mimimize the embedding dimension in order to reduce memory consumption. [Random projection](https://en.wikipedia.org/wiki/Random_projection) is a dimensionality reduction method for a set vectors in Euclidean space. This method is fast and requires no training. Let's try it with the model `EfficientNet-B7`, which generates embeddings of a high dimension 2560."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7ebf7b56",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A new collection created: tf_efficientnet_b7_512\n"
     ]
    }
   ],
   "source": [
    "NEW_MODEL = 'tf_efficientnet_b7'\n",
    "OLD_DIM = 2560\n",
    "NEW_DIM = 512\n",
    "NEW_COLLECTION_NAME = NEW_MODEL + '_' + str(NEW_DIM)\n",
    "\n",
    "numpy.random.seed(2023)\n",
    "projection_matrix = numpy.random.normal(scale=1.0, size=(OLD_DIM, NEW_DIM))\n",
    "\n",
    "def dim_reduce(vec):\n",
    "    return numpy.dot(vec, projection_matrix)\n",
    "\n",
    "connections.connect(host=HOST, port=PORT)\n",
    "new_collection = create_milvus_collection(NEW_COLLECTION_NAME, NEW_DIM)\n",
    "print(f'A new collection created: {NEW_COLLECTION_NAME}')\n",
    "\n",
    "\n",
    "# Embedding pipeline\n",
    "p_embed = (\n",
    "    pipe.input('img_path')\n",
    "        .map('img_path', 'img', ops.image_decode('rgb'))\n",
    "        .map('img', 'vec', ops.image_embedding.timm(model_name=NEW_MODEL, device=DEVICE))\n",
    "        .map('vec', 'vec', dim_reduce)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "23873a37",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table style=\"border-collapse: collapse;\"><tr><th style=\"text-align: center; font-size: 130%; border: none;\">img_path</th> <th style=\"text-align: center; font-size: 130%; border: none;\">img</th> <th style=\"text-align: center; font-size: 130%; border: none;\">vec</th></tr> <tr><td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \">./train/parachute/n03888257_31784.JPEG</td> <td style=\"text-align: center; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td> <td style=\"text-align: left; vertical-align: center; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \">[-3.1598568204094115, -6.1017211401300315, -16.71238753352162, ...] shape=(512,)</td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Display embedding result, no need for implementation\n",
    "p_display = p_embed.output('img_path', 'img', 'vec')\n",
    "\n",
    "DataCollection(p_display(to_insert[0])).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7519f65",
   "metadata": {},
   "source": [
    "We've build a new embedding pipeline converts each image into a vector of reduced dimension. Insert and search pipelines are built on top of the embedding pipeline, like what we did in previous sections. We can apply the same evaluation method to this engine.\n",
    "\n",
    "The dimension of embedding vector is reduced from 2560 to 512, thereby reducing memory usage by 80%. Despite this, it maintains a reasonable performance (97% mAP for reduced vectors vs 98.3% for full vectors)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "728dfe37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of data inserted: 1000\n",
      "mAP@10: 0.9700658698664651\n"
     ]
    }
   ],
   "source": [
    "# Insert pipeline\n",
    "p_insert = (\n",
    "        p_embed.map(('img_path', 'vec'), 'mr', ops.ann_insert.milvus_client(\n",
    "                    host=HOST,\n",
    "                    port=PORT,\n",
    "                    collection_name=NEW_COLLECTION_NAME\n",
    "                    ))\n",
    "          .output('mr')\n",
    ")\n",
    "\n",
    "# Insert data\n",
    "for img_path in to_insert:\n",
    "    p_insert(img_path)\n",
    "print('Number of data inserted:', new_collection.num_entities)\n",
    "\n",
    "# Search pipeline\n",
    "new_collection.load()\n",
    "p_search_pre = (\n",
    "        p_embed.map('vec', 'search_res', ops.ann_search.milvus_client(\n",
    "                    host=HOST, port=PORT, limit=TOPK,\n",
    "                    collection_name=NEW_COLLECTION_NAME))\n",
    "               .map('search_res', 'pred', lambda x: [str(Path(y[0]).resolve()) for y in x])\n",
    "#                .output('img_path', 'pred')\n",
    ")\n",
    "\n",
    "# Performance\n",
    "p_eval = (\n",
    "    p_search_pre.map('img_path', 'gt', ground_truth)\n",
    "                .map(('pred', 'gt'), 'ap', get_ap)\n",
    "                .output('ap')\n",
    ")\n",
    "\n",
    "res = []\n",
    "for img_path in to_test:\n",
    "    ap = p_eval(img_path).get()[0]\n",
    "    res.append(ap)\n",
    "\n",
    "mAP = mean(res)\n",
    "\n",
    "print(f'mAP@{TOPK}: {mAP}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb38ae46",
   "metadata": {},
   "source": [
    "## Online Demo <a class=\"anchor\" id=\"demo\"></a>\n",
    "\n",
    "This section shows how to use Gradio to build a simple showcase with user interface. With Gradio, we simply need to wrap the data processing pipeline via a f_search function. Please note here we search across a prepared Milvus collection *'deep_dive_image_search_vgg16'*, which stores normalized image embeddings extracted by vgg16."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b03bdefb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IMPORTANT: You are using gradio version 3.3, however version 3.14.0 is available, please upgrade.\n",
      "--------\n",
      "Running on local URL:  http://127.0.0.1:7860\n",
      "Running on public URL: https://c5d025b17a1dac43.gradio.app\n",
      "\n",
      "This share link expires in 72 hours. For free permanent hosting, check out Spaces: https://www.huggingface.co/spaces\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"https://c5d025b17a1dac43.gradio.app\" width=\"900\" height=\"500\" allow=\"autoplay; camera; microphone;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "(<gradio.routes.App at 0x7f45600d9280>,\n",
       " 'http://127.0.0.1:7860/',\n",
       " 'https://c5d025b17a1dac43.gradio.app')"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gradio\n",
    "\n",
    "DEMO_MODEL = 'vgg16'\n",
    "DEMO_COLLECTION = 'deep_dive_image_search_' + DEMO_MODEL\n",
    "\n",
    "def f_search(img):\n",
    "    p_search = (\n",
    "        pipe.input('img')\n",
    "            .map('img', 'vec', ops.image_embedding.timm(model_name=DEMO_MODEL, device=DEVICE))\n",
    "            .map('vec', 'vec', lambda x: x / numpy.linalg.norm(x, axis=0))\n",
    "            .map('vec', 'search_res', ops.ann_search.milvus_client(\n",
    "                    host=HOST, port=PORT, limit=TOPK,\n",
    "                    collection_name=DEMO_COLLECTION))\n",
    "            .map('search_res', 'pred', lambda x: [str(Path(y[0]).resolve()) for y in x])\n",
    "            .output('pred')\n",
    "    )\n",
    "    return p_search(img).get()[0]\n",
    "\n",
    "interface = gradio.Interface(f_search, \n",
    "                             gradio.inputs.Image(type=\"pil\", source='upload'),\n",
    "                             [gradio.outputs.Image(type=\"file\", label=None) for _ in range(TOPK)]\n",
    "                            )\n",
    "\n",
    "interface.launch(inline=True, share=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad827836",
   "metadata": {},
   "source": [
    "## Explore Towhee\n",
    "\n",
    "- Built-in pipelines for various tasks\n",
    "- Microservice & onnx acceleration powered by TritonServe\n",
    "- Docker image with everything ready"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28a5b088",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
